import re
from funasr import AutoModel


# paraformer-zh is a multi-functional asr model
# use vad, punc, spk or not as you need
model = AutoModel(model="paraformer-zh", model_revision="v2.0.4",
                  vad_model="fsmn-vad", vad_model_revision="v2.0.4",
                  punc_model="ct-punc-c", punc_model_revision="v2.0.4",
                  # spk_model="cam++", spk_model_revision="v2.0.2",
                  )


def only_asr(path):
    res = model.generate(input=path,
                         batch_size_s=300,
                         hotword='魔搭')

    texts = res[0]['text']
    print(texts)
    timestamp = res[0]['timestamp']
    merged_segments = []
    timestamp_len = len(timestamp)
    i = 0
    merged_text = ""
    start = 0
    end = 0
    for text in texts:
        print(i, text)
        # 标点符号和空格分段
        if re.match(r'^[，。、！？；,.!?\s]$', text):
            print("标点符号跳过", i)
            merged_segments.append({
                "text": merged_text,
                "start": start/1000,
                "end": end/1000
            })
            merged_text = ""
            start = 0
            end = 0
        else:
            if i >= timestamp_len:
                merged_text += text
                end = timestamp[timestamp_len-1][1]
                continue

            if merged_text == "":
                merged_text = text
                start = timestamp[i][0]
                end = timestamp[i][1]
            else:
                merged_text += text
                end = timestamp[i][1]
            i += 1
    return texts, merged_segments


if __name__ == "__main__":
    only_asr("1.wav")
